#!/usr/bin/python
'''
Pulls flanking sequences from a SAM file.
'''
import sys
import string
import fasta_read

# parameters.
sam_file = sys.argv[1]
fasta_file = sys.argv[2]
flank = int(sys.argv[3])

# read fasta.
seqs = fasta_read.read_fasta_dict(fasta_file)

# read sam.
fin = open(sam_file, "rb")
lines = fin.readlines()
fin.close()

# parse sam.
for line in lines:
	# tokenize.
	if line[0] == "@": continue
	tmp = line.split("\t")
	if tmp[2] == "*": continue

	# filter seq.
	dt = {}
	seq = tmp[9]
	for i in range(len(seq) - 2):
		dt[ seq[i:i+3] ] = True

	# check min.
	if len(dt) < 8: continue

	# check ref.
	if tmp[2] not in seqs: continue
	
	# pull out flank.
	bseq = seqs[tmp[2]]
	st = int(tmp[3])
	op = st + len(seq)

	# get bounds.
	if st - flank < 0:
		min = 0
	else:
		min = st - flank
	
	if op + flank >= len(bseq):
		max = len(bseq)
	else:
		max = op + flank

	fseq = bseq[min:max]
	fseq = fseq.upper()

	# flip reference if hit was mapped to other strand.
	if tmp[2] == "16":
		fseq.translate(string.maketrans("ATCGN", "TAGCN"))[::-1]

	# print info.
	print "%s\t%s\t%s" % (tmp[0], tmp[9], fseq)

